import sys
import io
import os
import json
import logging
import subprocess
import soundfile as sf
import numpy as np
import time
import warnings

# 修复Windows环境下的编码问题
if sys.platform.startswith('win'):
    # Windows环境下强制设置编码
    import codecs
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer)
    sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer)
    # 设置环境变量
    os.environ['PYTHONIOENCODING'] = 'utf-8'
else:
    # 非Windows环境使用原有方式
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

def safe_print(message):
    """安全的打印函数，处理编码问题"""
    try:
        print(message, file=sys.stderr)
    except UnicodeEncodeError:
        # 如果UTF-8编码失败，尝试使用系统默认编码
        print(message.encode('utf-8', errors='ignore').decode('utf-8', errors='ignore'), file=sys.stderr)

def check_model_availability():
    """检查模型可用性"""
    models = {}
    
    # 检查Whisper
    try:
        import whisper
        models['whisper'] = {
            'available': True,
            'name': 'Whisper',
            'description': '高精度识别，适合高质量要求',
            'speed': '中等',
            'accuracy': '高'
        }
    except ImportError:
        models['whisper'] = {
            'available': False,
            'name': 'Whisper',
            'description': '未安装',
            'speed': 'N/A',
            'accuracy': 'N/A'
        }
    
    # 检查Vosk
    try:
        import vosk
        models['vosk'] = {
            'available': True,
            'name': 'Vosk',
            'description': '快速识别，适合实时应用',
            'speed': '快',
            'accuracy': '中等'
        }
    except ImportError:
        models['vosk'] = {
            'available': False,
            'name': 'Vosk',
            'description': '未安装',
            'speed': 'N/A',
            'accuracy': 'N/A'
        }
    
    return models

def auto_select_model(audio_file_path, preference='auto'):
    """自动选择最佳模型"""
    models = check_model_availability()
    available_models = {k: v for k, v in models.items() if v['available']}
    
    if not available_models:
        return None, "没有可用的语音识别模型"
    
    # 根据偏好选择模型
    if preference == 'accuracy':
        # 优先选择高精度模型
        if 'whisper' in available_models:
            return 'whisper', "选择Whisper：高精度识别"
        elif 'vosk' in available_models:
            return 'vosk', "选择Vosk：可用模型"
    
    elif preference == 'speed':
        # 优先选择快速模型
        if 'vosk' in available_models:
            return 'vosk', "选择Vosk：快速识别"
        elif 'whisper' in available_models:
            return 'whisper', "选择Whisper：可用模型"
    
    else:  # preference == 'auto'
        # 自动选择：根据音频长度和文件大小
        try:
            audio_info = sf.info(audio_file_path)
            duration = audio_info.duration
            file_size = os.path.getsize(audio_file_path) / (1024 * 1024)  # MB
            
            if duration < 10 and file_size < 5:  # 短音频，选择快速模型
                if 'vosk' in available_models:
                    return 'vosk', f"选择Vosk：短音频({duration:.1f}s, {file_size:.1f}MB)，快速识别"
                else:
                    return 'whisper', f"选择Whisper：短音频，可用模型"
            else:  # 长音频，选择高精度模型
                if 'whisper' in available_models:
                    return 'whisper', f"选择Whisper：长音频({duration:.1f}s, {file_size:.1f}MB)，高精度识别"
                else:
                    return 'vosk', f"选择Vosk：长音频，可用模型"
        except:
            # 如果无法获取音频信息，默认选择Whisper
            if 'whisper' in available_models:
                return 'whisper', "选择Whisper：默认选择"
            else:
                return 'vosk', "选择Vosk：默认选择"
    
    return None, "模型选择失败"

def recognize_with_whisper(audio_file_path):
    """使用Whisper进行语音识别"""
    try:
        import whisper
        
        safe_print("🔄 正在加载Whisper模型...")
        model = whisper.load_model("base")
        
        safe_print("🎤 正在进行Whisper语音识别...")
        start_time = time.time()
        
        result = model.transcribe(audio_file_path, language="zh")
        
        end_time = time.time()
        processing_time = end_time - start_time
        
        safe_print(f"✅ Whisper识别完成！耗时: {processing_time:.2f}秒")
        
        return {
            "success": True,
            "text": result["text"],
            "confidence": 0.9,  # Whisper不提供置信度，设为默认值
            "processing_time": processing_time,
            "model": "Whisper",
            "language": "中文"
        }
        
    except Exception as e:
        safe_print(f"❌ Whisper识别失败: {e}")
        return {"error": f"Whisper识别失败: {e}"}

def recognize_with_vosk(audio_file_path):
    """使用Vosk进行语音识别"""
    try:
        from vosk import Model, KaldiRecognizer
        
        # 检查是否有中文模型
        model_path = "vosk-model-small-cn-0.22"
        if not os.path.exists(model_path):
            safe_print("⚠️ 未找到Vosk中文模型，使用默认模型")
            model = Model(lang="en-us")
        else:
            safe_print(f"🔄 正在加载Vosk中文模型: {model_path}")
            model = Model(model_path)
        
        safe_print("🎤 正在进行Vosk语音识别...")
        start_time = time.time()
        
        # 读取音频
        audio_data, sample_rate = sf.read(audio_file_path)
        if len(audio_data.shape) > 1:
            audio_data = audio_data[:, 0]
        
        # 转换为16kHz
        if sample_rate != 16000:
            if sample_rate > 16000:
                ratio = sample_rate // 16000
                audio_data = audio_data[::ratio]
            else:
                ratio = 16000 // sample_rate
                audio_data = np.repeat(audio_data, ratio)
            sample_rate = 16000
        
        # 识别
        rec = KaldiRecognizer(model, sample_rate)
        audio_bytes = (audio_data * 32767).astype(np.int16).tobytes()
        rec.AcceptWaveform(audio_bytes)
        result = rec.FinalResult()
        
        end_time = time.time()
        processing_time = end_time - start_time
        
        # 解析结果
        result_json = json.loads(result)
        text = result_json.get("text", "")
        confidence = result_json.get("confidence", 0.0)
        
        safe_print(f"✅ Vosk识别完成！耗时: {processing_time:.2f}秒")
        
        return {
            "success": True,
            "text": text,
            "confidence": confidence,
            "processing_time": processing_time,
            "model": "Vosk",
            "language": "中文" if "vosk-model-small-cn" in model_path else "英文"
        }
        
    except Exception as e:
        safe_print(f"❌ Vosk识别失败: {e}")
        return {"error": f"Vosk识别失败: {e}"}

def main():
    """主函数"""
    if len(sys.argv) < 2:
        safe_print("❌ 用法: python speech_recognize_smart.py <音频文件路径> [preference]")
        safe_print("    preference: auto(默认), accuracy(高精度), speed(快速)")
        sys.exit(1)
    
    audio_file_path = sys.argv[1]
    preference = sys.argv[2] if len(sys.argv) > 2 else 'auto'
    
    if not os.path.exists(audio_file_path):
        safe_print(f"❌ 音频文件不存在: {audio_file_path}")
        sys.exit(1)
    
    safe_print("🎤 开始智能语音识别...")
    safe_print("=" * 50)
    
    # 显示可用模型
    models = check_model_availability()
    safe_print("📋 可用模型:")
    for model_id, model_info in models.items():
        status = "✅ 可用" if model_info['available'] else "❌ 不可用"
        safe_print(f"  {model_info['name']}: {status} - {model_info['description']}")
    
    safe_print("=" * 50)
    
    # 自动选择模型
    selected_model, reason = auto_select_model(audio_file_path, preference)
    if not selected_model:
        safe_print(f"❌ {reason}")
        sys.exit(1)
    
    safe_print(f"🤖 {reason}")
    safe_print(f"🎯 选择模型: {models[selected_model]['name']}")
    
    # 使用选定的模型进行识别
    if selected_model == 'whisper':
        result = recognize_with_whisper(audio_file_path)
    elif selected_model == 'vosk':
        result = recognize_with_vosk(audio_file_path)
    else:
        safe_print(f"❌ 未知模型: {selected_model}")
        sys.exit(1)
    
    # 输出结果
    if result.get("success"):
        safe_print("=" * 50)
        safe_print("🎉 识别成功！")
        safe_print(f"📝 文本: {result['text']}")
        safe_print(f"🎯 置信度: {result['confidence']:.2f}")
        safe_print(f"⏱️ 处理时间: {result['processing_time']:.2f}秒")
        safe_print(f"🤖 模型: {result['model']}")
        safe_print(f"🌍 语言: {result['language']}")
        
        # 输出JSON格式结果 - 这是关键，确保输出正确的JSON格式
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        safe_print("=" * 50)
        safe_print("❌ 识别失败！")
        safe_print(f"错误: {result.get('error', '未知错误')}")
        # 即使失败也要输出JSON格式
        print(json.dumps(result, ensure_ascii=False, indent=2))
        sys.exit(1)

if __name__ == "__main__":
    main()
